"""ChatBot implementation for WeChat integration"""
from llm import Chat
from wxauto import WeChat
import os
import time
import random
from collections import defaultdict
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('chat.log', encoding='utf-8')
    ]
)

class ChatBot:
    def __init__(self, api_key, system_prompt, listen_list, merge_wait_time=2.0):
        # 初始化DeepSeek实例
        self.deepseek = Chat(api_key=api_key, prompt=system_prompt)
        # 初始化微信实例
        self.wx = WeChat()
        # 保存配置
        self.listen_list = listen_list
        # 初始化消息缓冲区
        self.message_buffer = defaultdict(list)
        # 初始化时间记录
        self.last_message_time = defaultdict(float)
        # 设置合并等待时间
        self.merge_wait_time = merge_wait_time
        # 初始化聊天窗口缓存
        self.chat_windows = {}
        
        # 初始化监听列表
        for name in listen_list:
            try:
                self.wx.AddListenChat(who=name)
                logging.info(f"添加监听: {name}")
            except Exception as e:
                logging.error(f"添加监听失败 - {name}: {str(e)}")

    def process_new_messages(self):
        """处理新消息"""
        try:
            msgs = self.wx.GetListenMessage()
            current_time = time.time()
            
            # 更新聊天窗口缓存
            for chat in msgs:
                msg = msgs.get(chat)
                self.chat_windows[chat] = msg
                
                for item in msg:
                    if item.type == 'friend':
                        self.message_buffer[item.sender].append(item.content)
                        self.last_message_time[item.sender] = current_time
                        logging.info(f"收到消息 - {item.sender}: {item.content}")
            
            return current_time
        except Exception as e:
            logging.error(f"处理新消息时出错: {str(e)}")
            return time.time()

    def send_reply(self, sender, reply):
        """发送回复消息"""
        try:
            for chat in self.chat_windows:
                msg = self.chat_windows.get(chat)
                for item in msg:
                    if item.type == 'friend' and item.sender == sender:
                        chat.SendMsg(reply)
                        logging.info(f"发送消息 - {sender}: {reply}")
                        return True
            return False
        except Exception as e:
            logging.error(f"发送回复时出错 - {sender}: {str(e)}")
            return False

    def process_buffered_messages(self, current_time):
        """处理缓冲区中的消息"""
        for sender in list(self.message_buffer.keys()):
            try:
                if (current_time - self.last_message_time[sender] >= self.merge_wait_time 
                    and self.message_buffer[sender]):
                    
                    combined_message = "\n".join(self.message_buffer[sender])
                    logging.info(f"处理合并消息 - {sender}:\n{combined_message}")
                    
                    reply = self.deepseek.chat(combined_message, user_id=sender)
                    time.sleep(random.uniform(1, 3))  # 随机等待1-3秒
                    
                    if self.send_reply(sender, reply):
                        self.message_buffer[sender] = []  # 清空已处理的消息
                    
            except Exception as e:
                logging.error(f"处理缓冲消息时出错 - {sender}: {str(e)}")

    def run(self):
        """运行聊天机器人"""
        print("ChatBot已启动，按Ctrl+C停止运行")
        while True:
            current_time = self.process_new_messages()
            self.process_buffered_messages(current_time)
            time.sleep(0.5)  # 短暂休眠以减少CPU使用